import numpy as np
import torch
import torchvision.transforms as transforms
import cv2
import glob
import os
import random
from PIL import Image,ImageOps

class LolTrainLoader(torch.utils.data.Dataset):
    
    def __init__(self, gt_dir,img_dir,patch_size,mode = 'train',he_ops=False):
        self.low_img_dir = img_dir
        self.gt_dir = gt_dir
        self.he_ops=he_ops
        self.patch_size = patch_size
        self.mode = mode
    
        self.low_data_names = _load_one_dir(img_dir)
     
        self.low_data_names.sort()
        self.gt_data_names=_load_one_dir(gt_dir)
        self.gt_data_names.sort()

        self.count = len(self.low_data_names)
        transform_list = []
        transform_list += [transforms.ToTensor()]
        self.transform = transforms.Compose(transform_list)

    def load_images_transform(self, file,he_ops=False):
        im = Image.open(file).convert('RGB')
        if he_ops:
            im = ImageOps.equalize(im)

        img_norm = self.transform(im).numpy()
        img_norm = np.transpose(img_norm, (1, 2, 0))
        return img_norm

    def __getitem__(self, index):
        img_name=self.low_data_names[index].split('/')[-1]
        low = self.load_images_transform(self.low_data_names[index],self.he_ops)
     
        high = self.load_images_transform(self.gt_data_names[index])
        if self.mode=='train':
            h = low.shape[0]
            w = low.shape[1]

            h_offset = random.randint(0, max(0, h - self.patch_size - 1))
            w_offset = random.randint(0, max(0, w - self.patch_size - 1))

            # crop
            low = low[h_offset:h_offset + self.patch_size, w_offset:w_offset + self.patch_size]
            high = high[h_offset:h_offset + self.patch_size, w_offset:w_offset + self.patch_size]

        
        low = np.asarray(low, dtype=np.float32)
        low = np.transpose(low[:, :, :], (2, 0, 1))
        high = np.asarray(high, dtype=np.float32)
        high = np.transpose(high[:, :, :], (2, 0, 1))


        return torch.from_numpy(low), torch.from_numpy(high), img_name

    def __len__(self):
        return self.count